#include "library.h"

DWORD post = 0xE61AD1;

void __declspec(naked) Hook() {

    __asm {

        call MessageBoxA
        jmp post
    }

}

BOOL WINAPI SetHardwareBP(DWORD address, HANDLE thread) {

    // Win10下必须先挂起线程再设置线程上下文
    SuspendThread(thread);

    // 构建线程上下文对象并设置断点
    CONTEXT ctx = { CONTEXT_DEBUG_REGISTERS };          // 创建线程Context对象
    ctx.Dr7 = 0x1;                                      // 启用Dr0
    ctx.Dr0 = address;                                  // 设置硬件断点

    // 应用线程上下文
    BOOL res = SetThreadContext(thread, &ctx);

    // 恢复线程运行
    ResumeThread(thread);

    return res;

}

BOOL WINAPI SetVehHook(DWORD hookAddr) {

    UINT pid = GetCurrentProcessId();

    // 0为当前进程
    HANDLE threads = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, pid);
    if (threads == INVALID_HANDLE_VALUE) return FALSE;

    // 线程环境结构体
    THREADENTRY32 curThreadEntry = { 0 };
    curThreadEntry.dwSize = sizeof(THREADENTRY32);

    // 遍历中的当前线程
    HANDLE curThread;

    // 遍历当前进程的所有线程
    if (Thread32First(threads, &curThreadEntry)) {
        do {

            // 判断是否为当前进程的线程
            if (curThreadEntry.th32OwnerProcessID == pid) {

                // 获取线程句柄
                curThread = OpenThread(THREAD_SET_CONTEXT | THREAD_GET_CONTEXT | THREAD_QUERY_INFORMATION, FALSE, curThreadEntry.th32ThreadID);
                if (curThread == INVALID_HANDLE_VALUE) continue;

                // 设置硬件断点
                SetHardwareBP(hookAddr, curThread);

                CloseHandle(curThread);
            }
        } while (Thread32Next(threads, &curThreadEntry));
    }

    CloseHandle(threads);

    return TRUE;
}

// 异常处理函数
LONG NTAPI ExceptionFilter(PEXCEPTION_POINTERS ExceptionInfo) {

    // 异常是否为硬件断点触发的单步执行异常
    if (ExceptionInfo->ExceptionRecord->ExceptionCode == EXCEPTION_SINGLE_STEP) {

        // 触发异常的地址是否是我们需要hook的地址
        if (reinterpret_cast<DWORD>(ExceptionInfo->ExceptionRecord->ExceptionAddress) == 0xE61ACB) {

            MessageBoxA(nullptr, "Veh Hook Succeeded!", "Message from Ayin", MB_OK);

            PCONTEXT context = ExceptionInfo->ContextRecord;
            context->Eip = reinterpret_cast<DWORD>(Hook);

            return EXCEPTION_CONTINUE_EXECUTION;
        }
    }

    return EXCEPTION_CONTINUE_SEARCH;
}



